#!/usr/bin/env python3

import json
import logging
import os
import stat
import subprocess

from argparse import ArgumentParser

ESP_GUID = 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B'
MNT_DIR = '/mnt'

logger = logging.getLogger(__name__)


def path_fixup():
    path = os.environ.get('PATH', '/bin:/usr/bin:/sbin:/usr/sbin')
    path = path.split(':')
    for thedir in ('/sbin', '/usr/sbin'):
        if thedir not in path:
            path.append(thedir)
    path = ':'.join(path)
    os.environ['PATH'] = path


def parent_blockdev(devnode):
    part = devnode
    if part.startswith('/dev/'):
        part = devnode[5:]
    parent_sysfs = f'/sys/class/block/{part}/..'
    resolved = os.path.realpath(parent_sysfs)
    dev = os.path.basename(resolved)
    return dev


def read_partition_table(drive):
    cmd = ['sfdisk', '-l', '--json', drive]
    logger.debug('read_partition_table: running "%s"', ' '.join(cmd))
    out = subprocess.check_output(cmd, encoding='utf-8')
    return json.loads(out.strip())['partitiontable']['partitions']


def find_efi_syspart(drive):
    partitions = read_partition_table(drive)
    esp_candidates = [p for p in partitions if p['type'] == ESP_GUID]
    logger.debug('find_efi_syspart: got %d ESP candidates', len(esp_candidates))
    if len(esp_candidates) == 1:
        return esp_candidates[0]
    elif len(esp_candidates) == 0:
        raise RuntimeError("Couldn't find ESP on %s" % drive)
    else:
        raise RuntimeError('Several (%d) ESPs found on %s',
                           len(esp_candidates), drive)


def verify_dest_root(devnode):
    st = os.stat(devnode)
    if not stat.S_ISBLK(st.st_mode):
        logger.warning('%s is not a block device', devnode)
        return False


def run_mke2fs(dev):
    cmd = 'mke2fs -t ext4 -b 4096 -L BoringLinux'.split()
    cmd += [dev]
    logger.debug('formatting rootfs: %s', ' '.join(cmd))
    subprocess.check_call(cmd)


def mount_rootfs(dev):
    cmd = 'mount -t ext4 -o rw,data=ordered,noatime,errors=remount-ro'.split()
    cmd += [dev, MNT_DIR]
    logger.debug('mounting rootfs: %s', ' '.join(cmd))
    subprocess.check_call(cmd)


def untar(tarball, destdir):
    cmd = 'tar -x -a -v -p --numeric-owner -f'.split()
    cmd += [tarball, '-C', destdir]
    logger.debug('unpacking rootfs: %s', ' '.join(cmd))
    subprocess.check_call(cmd)


def mount_efi_syspart(dev):
    cmd = 'mount -t vfat -o rw,codepage=866,iocharset=ascii,utf8,noatime'
    cmd = cmd.split()
    cmd += [dev, f'{MNT_DIR}/boot/efi']
    logger.debug('mounting EFI system partition: %s', ' '.join(cmd))
    subprocess.check_call(cmd)


def install_grub(boot_dir, efi_dir, removable=False):
    cmd = [
        'grub-install',
        '--boot-directory={boot_dir}',
        '--efi-directory={efi_dir}',
    ]
    if removable:
        cmd += [
            '--no-nvram',
            '--removable'
        ]
    else:
        cmd += [
            '--bootloader-id=boringlinux'
        ]
    logger.debug('installing grub: %s', ' '.join(cmd))

    subprocess.check_call(cmd)


def enable_service(unit, chroot):
    cmd = [
        'systemctl',
        'enable',
        f'--root={chroot}',
        unit
    ]
    logger.debug('enabling %s: %s', unit, ' '.join(cmd))
    subprocess.check_call(cmd)


def umount(mntpt, *args):
    cmd = ['umount', mntpt]
    cmd += list(*args)
    logger.debug('unmounting %s: %s', mntpt, ' '.join(cmd))
    subprocess.check_call(cmd)


def main():
    logging.basicConfig(format='%(asctime)-15s %(levelname)s: %(message)s',
                        level=logging.DEBUG)
    path_fixup()
    parser = ArgumentParser('Install Boring Linux from tarball')
    parser.add_argument('-e', '--esp', help='Destination EFI system partition')
    parser.add_argument('-r', '--removable',
                        help='Install GRUB into EFI/BOOT')
    parser.add_argument('src', help='Boring Linux installation tarball')
    parser.add_argument('dest', help='Destination root filesystem')
    args = parser.parse_args()
    print(args.dest)
    is_blk = verify_dest_root(args.dest)
    if is_blk:
        disk = parent_blockdev(args.dest)
    else:
        disk = args.dest
    logger.info('parent drive: %s', disk)
    esp = find_efi_syspart(disk)
    run_mke2fs(args.dest)
    mount_rootfs(args.dest)
    unpack_tarball(args.src, MNT_DIR)
    for service in ('systemd-networkd', 'systemd-resolved'):
        enable_service(service, MNT_DIR)
    mount_efi_syspart(esp)
    install_grub(f'{MNT_DIR}/boot', f'{MNT_DIR}/boot/efi', args.removable)
    umount(f'{MNT_DIR}/boot/efi')


if __name__ == '__main__':
    main()

